Codecov Report

Attention: Patch coverage is 60.29777% with 160 lines in your changes missing coverage. Please review.

Thanks @vpratz for the implementation! sigma_data is already specific for the EDM version. I think, we have to make it even more general, so some additional arguments which can be passed, depending on the type of schedule one wants to use.

I plan to add on top of this additional schedules and samplers until the end of the week.

Thanks for taking a look. Do you know whether your implementation would benefit from the pre-conditioning discussed in Elucidating the Design Space of Diffusion-Based Generative Models, and whether we can combine them in one joint framework?

Part of the pre-conditioning can be expressed as a special kind of weighting function: see appendix D.1 in here.

So yes, the aim would be to have one nice framework!

I added some more noise schedules and started to make the implementation more general. This is just a first draft, so you @vpratz get an idea, how we could do it. We should discuss this then and how to move forward.

I added a class NoiseSchedule and different schedules, so it should be easy now to extend to more schedules if necessary. Since EDM has a specific sampling scheme for inference, this is now also defined in the noise schedule. Therefore, we do not have to specify specific sampling step sizes anymore.

Next step would be add stochastic samplers as well.

Thanks a lot for the fixes, they increase the performance of the new implementation a lot. The old standalone EDM implementation seems to be a little bit better still, but the difference might be down to hyperparameter tuning. I have added the examples/experimental/Two_Moons_Diffusion_Comparison.ipynb notebook, which allows testing both implementations on the same benchmark, and plotting the results against each other. This is for two moons, but feel free to expand it with other benchmarks as well.

As far as I can tell, the open steps before we finalize this PR are:

• optimising performance: ensuring that we achieve the same performance as with the standalone EDM implementation, so that we do not forego performance by not including it
• related: find and set good defaults
• remove the LinearNoiseSchedule (as discussed privately, maybe moving it to a tutorial): the other schedules perform better, so we do not need to include and maintain it here
• add tests to cover all relevant cases/combinations
• proof-reading docstrings and maybe supplying a tutorial
• remove the examples/experimental/Two_Moons_Diffusion_Comparison.ipynb notebook before merging

Did I miss anything, @arrjon , or do you have any other comments on the current state?

The performance issue is fixed now, it was mainly due to a missing scaling factor of the log_snr, which goes into the network.

I am also implementing the stochastic sampler: it is working for all backend but jax at the moment. After this, only the things @vpratz mentioned are missing.

The stochastic sampler is now also working for jax. So all features done for the moment!

Great! Thanks a lot for putting in the work and for the quick fixes! I'll try to add the relevant tests and work on some of the other missing things in the next few days.